import json

def get_source(tokens,new_function_name=None,use_identify=False):
    def identify(token):
        if token in ['#INDENT#','#NEWLINE#']:
            return 1
        elif token.startswith('def '):
            return 0
        else:
            return -1  
    from io import StringIO
    buffer=StringIO()
    
    start=0
    active=False
    
    function_name_index=None
    
    if new_function_name:
        for i,token in enumerate(tokens):
            if token.startswith("def "):
                function_name_index=i
                break
        
    for i,token in enumerate(tokens):
        if use_identify:
            status=identify(token)
            if status==0:
                use_identify=False
            elif status==1:
                pass
            else:
                continue
            
        if function_name_index==i:
            token=new_function_name
            
        if token=="#INDENT#":
            start+=1
        elif token=='#UNINDENT#':
            start-=1
        elif token=='#NEWLINE#':
            buffer.write('\n')
            active=True
        else:
            start_indent="\t"*start
            if active:
                buffer.write(start_indent+token)       
                active=False
            else:
                buffer.write(" "+token)
    buffer.write('\n')
    s=buffer.getvalue()
    buffer.close()
    return s           

class Py150kDataAdatper:
    
    @staticmethod
    def line2data(json_data):
        def parse_edge(edges,upper_bound):
            result={}
            for edge in edges:
                if max(edge[:2])>=upper_bound or min(edge[:2])<=0: continue
                if edge[2] not in result:
                    result[edge[2]]=[]
                result[edge[2]].append(edge[:2])
            return result
        json_data=json.loads(json_data)
        #print(json_data)
        tokens=json_data["source_tokens"]
        error_location=json_data["error_location"]
        repair_candidates=[x for x in json_data["repair_candidates"] if isinstance(x,int)]
        has_bug=json_data["has_bug"]
        repair_targets=json_data["repair_targets"]
        edges=json_data["edges"]
        
        edges=parse_edge(edges,len(tokens))
        
        source=get_source(tokens)
            
        return source,tokens,error_location,repair_candidates,repair_targets,edges,has_bug
    
    @staticmethod
    def from_varmisuse_dataset(path):

        with open(path,'r') as f:
            for l in f:
                yield Py150kDataAdatper.line2data(l)

    @staticmethod
    def iter_varmisuse_dataset(path):
        for e in Py150kDataAdatper.from_varmisuse_dataset(path):
            source,tokens,error_location,repair_candidates,repair_targets,edges,has_bug=e
            yield has_bug,tokens,source,error_location,repair_candidates,repair_targets,edges